import io
import os
import time

import airtest.core.android.adb
from PIL import Image
from airtest.core.api import snapshot


def get_screenshot_formac(adb: airtest.core.android.adb.ADB,
                          save_dir: str,
                          filename_format: str = "screenshot_%Y%m%d_%H%M%S.png") -> str:
    """
    使用 Airtest 的 ADB 模块进行截图，并保存到指定目录。

    :param filename_format: 文件名格式，支持时间格式化字符串（如 %Y%m%d_%H%M%S）。
    :param save_dir: 保存截图的目录路径。
    :param adb:

    :return 生成的截图文件路径。
    """
    try:
        # 确保目标目录存在
        os.makedirs(save_dir, exist_ok=True)

        # 生成带时间戳的文件名
        timestamp = time.strftime(filename_format)
        filename = os.path.join(save_dir, timestamp)

        # 截图
        image_array = adb.snapshot()
        image = Image.open(io.BytesIO(image_array))
        image.save(filename, format="png", quality=85)
        return filename

    except Exception as e:
        print(f"[ERROR] 截图失败: {e}")
        raise


def get_screenshot(save_dir: str,
                   filename_format: str = "screenshot_%Y%m%d_%H%M%S.png") -> str:
    """
    使用 Airtest 的 ADB 模块进行截图，并保存到指定目录。

    :param filename_format: 文件名格式，支持时间格式化字符串（如 %Y%m%d_%H%M%S）。
    :param save_dir: 保存截图的目录路径。
    :param adb:

    :return 生成的截图文件路径。
    """

    # 确保目标目录存在
    os.makedirs(save_dir, exist_ok=True)

    # 生成带时间戳的文件名
    timestamp = time.strftime(filename_format)
    filename = os.path.join(save_dir, timestamp)

    # 截图
    snapshot(filename)

    return filename
